package ffte.core.flat

import spinal.core._
import spinal.lib._

import ffte.types.VecComplex
import ffte.property.{getFFTIW,getFFTOW,FFTIOConfig}
import ffte.algorithm.FFTGen.{Method,psbFFT}

/**
 * Prime-Space Block (PSB) FFT implementation
 *
 * This class implements FFT using the Prime-Space Block decomposition method,
 * which is based on the Chinese Remainder Theorem. PSB decomposition is particularly
 * efficient for transform sizes that can be factored into coprime numbers,
 * enabling optimal hardware mapping through index remapping.
 *
 * Mathematical Foundation:
 * The PSB method decomposes an N-point DFT where N = N1 * N2 with gcd(N1, N2) = 1.
 * Using the Chinese Remainder Theorem, the input and output indices can be
 * remapped to separate the transform into two independent stages:
 *
 * Index Mapping:
 * - Input: n = (n1*N2*inv_N2_mod_N1 + n2*N1*inv_N1_mod_N2) mod N
 * - Output: k = (k1*N2*inv_N2_mod_N1 + k2*N1*inv_N1_mod_N2) mod N
 *
 * Where inv_N2_mod_N1 is the modular inverse of N2 modulo N1.
 *
 * Hardware Architecture:
 * - First stage: N2 independent N1-point FFTs (parallel computation)
 * - Second stage: N1 independent N2-point FFTs (parallel computation)
 * - Interconnection: Permutation network based on CRT mapping
 * - No twiddle factors required due to coprime decomposition
 *
 * Performance Characteristics:
 * - Optimal for coprime factorizations
 * - Minimal multiplication count (no twiddle factors)
 * - High parallelism through independent FFT blocks
 * - Regular memory access patterns
 *
 * @tparam T Method type for sub-FFT implementations (Flat, Mixed, Streamed)
 * @param S Scaling shift amount for fixed-point overflow prevention
 * @param p PSB decomposition structure containing stage information and mapping tables
 */
class psb[T <: Method](val S: Int, val p: psbFFT[T]) extends combineFlatComponents[T] {
    /** Standard I/O interface for flat FFT components */
    val io = iogen

    /**
     * First stage FFT processing block
     *
     * This block implements the first stage of the PSB decomposition, consisting
     * of N2 independent N1-point FFTs operating in parallel. Each FFT processes
     * a subset of the input data according to the Chinese Remainder Theorem mapping.
     *
     * Architecture:
     * - N2 parallel FFT processors of size N1
     * - Input width: dW (global input precision)
     * - Output width: interW (intermediate precision)
     * - All processors share the same clock enable signal
     *
     * Hardware Mapping:
     * Each processor handles inputs at indices congruent to a specific value
     * modulo N2, implementing the first dimension of the CRT decomposition.
     */
    val f_block = FFTIOConfig(dW, interW) on new Area {
        /**
         * Array of first-stage FFT processors
         *
         * Creates N2 independent FFT processors, each handling N1 points.
         * The processors are named "f0", "f1", ..., "f(N2-1)" for debugging.
         */
        val u = (0 until second.N).toList.map { i => fg(first, s"f$i") }

        /**
         * Delay through first stage
         *
         * All processors have identical delay since they use the same
         * implementation method and transform size.
         */
        val delay = u(0).delay

        /**
         * Clock enable distribution
         *
         * All first-stage processors share the same clock enable signal,
         * ensuring synchronous operation across the parallel array.
         */
        u.foreach { x => x.io.ce := io.ce }
    }

    /**
     * Second stage FFT processing block
     *
     * This block implements the second stage of the PSB decomposition, consisting
     * of N1 independent N2-point FFTs operating in parallel. Each FFT processes
     * the intermediate results from the first stage according to the CRT mapping.
     *
     * Architecture:
     * - N1 parallel FFT processors of size N2
     * - Input width: interW (intermediate precision from first stage)
     * - Output width: oW (global output precision)
     * - All processors share the same clock enable signal
     *
     * Hardware Mapping:
     * Each processor handles intermediate data corresponding to a specific
     * residue class modulo N1, completing the two-dimensional CRT decomposition.
     */
    val s_block = FFTIOConfig(interW, oW) on new Area {
        /**
         * Array of second-stage FFT processors
         *
         * Creates N1 independent FFT processors, each handling N2 points.
         * The processors are named "f0", "f1", ..., "f(N1-1)" for debugging.
         */
        val u = (0 until first.N).toList.map { i => sg(second, s"f$i") }

        /**
         * Delay through second stage
         *
         * All processors have identical delay since they use the same
         * implementation method and transform size.
         */
        val delay = u(0).delay

        /**
         * Clock enable distribution
         *
         * All second-stage processors share the same clock enable signal,
         * ensuring synchronous operation across the parallel array.
         */
        u.foreach { x => x.io.ce := io.ce }
    }

    /**
     * Data interconnection network
     *
     * Implements the Chinese Remainder Theorem-based data routing between
     * the first and second stage FFT processors. This network performs the
     * index remapping required for PSB decomposition.
     *
     * Connection Pattern:
     * - Input to first stage: Direct assignment based on row-major ordering
     * - First to second stage: CRT-based transposition
     * - Second stage to output: Inverse CRT mapping
     *
     * Index Mathematics:
     * - Input index: i*first.N + k maps to processor i, sample k
     * - Intermediate: processor i, sample k maps to processor k, sample i
     * - Output index: k*second.N + i maps from processor k, sample i
     */
    for (i <- 0 until second.N) {
        for (k <- 0 until first.N) {
            /**
             * First stage to second stage connection
             *
             * Routes the output of first-stage processor i, sample k
             * to the input of second-stage processor k, sample i.
             * This implements the transpose operation required by CRT.
             */
            s_block.u(k).io.d.d(i) := f_block.u(i).io.q.d(k)

            /**
             * Input to first stage connection
             *
             * Routes global input data to the appropriate first-stage processor
             * based on row-major ordering of the 2D decomposition.
             */
            f_block.u(i).io.d.d(k) := io.d.d(i * first.N + k)

            /**
             * Second stage to output connection
             *
             * Routes the output of second-stage processor k, sample i
             * to the global output using column-major ordering.
             */
            io.q.d(k * second.N + i) := s_block.u(k).io.q.d(i)
        }
    }

    /**
     * Total component delay in clock cycles
     *
     * The total delay is the sum of delays through both stages since
     * they operate sequentially in the data flow path.
     *
     * @return Total delay from input to output in clock cycles
     */
    def delay = f_block.delay + s_block.delay
}
